import cv2
from ultralytics import YOLO
import torch
import numpy as np

class YoloPredictor:
    """
    一个YOLO预测器类，用于加载模型、执行预测、跟踪和实例分割。
    增加了对红色和橙色标签的颜色校验逻辑。
    """
    def __init__(self, model_path):
        """
        初始化预测器，加载YOLO模型。
        :param model_path: YOLO模型文件的路径（.pt文件）。
        """
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"正在使用设备: {self.device}")
        self.model = YOLO(model_path)
        # self.model.to(self.device)
        self.results = None
        self.current_frame = None

    def predict_and_track(self, frame):
        """
        对单帧图像执行目标检测、实例分割和跟踪。
        :param frame: 从摄像头捕获的图像帧。
        :return: Ultralytics的Results对象。
        """
        self.current_frame = frame  # 存储当前帧以供颜色分析
        self.results = self.model.predict(frame, verbose=False,show=True)
        return self.results

    def _refine_red_orange_classification(self, box_coords):
        """
        通过分析HSV颜色空间中的像素来精确区分红色和橙色。
        这是一个内部辅助方法。
        :param box_coords: 物体的边界框坐标 [x1, y1, x2, y2]。
        :return: 'red'、'orange' 或 'unknown' 字符串。
        """
        x1, y1, x2, y2 = map(int, box_coords)
        
        # 确保坐标在图像范围内
        h, w, _ = self.current_frame.shape
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(w, x2), min(h, y2)

        cropped_object = self.current_frame[y1:y2, x1:x2]

        if cropped_object.size == 0:
            return "unknown"

        hsv_image = cv2.cvtColor(cropped_object, cv2.COLOR_BGR2HSV)

        # --- HSV 颜色范围定义 ---
        # 注意：这些值可能需要根据您的具体光照条件和摄像头进行微调。
        # 红色在HSV中环绕0/180，所以需要两个范围。
        lower_red1 = np.array([0, 120, 70])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([170, 120, 70])
        upper_red2 = np.array([180, 255, 255])

        # 橙色范围
        lower_orange = np.array([11, 120, 70])
        upper_orange = np.array([25, 255, 255])
        # ---

        # 创建颜色掩码
        mask_red1 = cv2.inRange(hsv_image, lower_red1, upper_red1)
        mask_red2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
        red_mask = cv2.add(mask_red1, mask_red2)
        orange_mask = cv2.inRange(hsv_image, lower_orange, upper_orange)

        # 计算颜色像素数量
        red_pixel_count = cv2.countNonZero(red_mask)
        orange_pixel_count = cv2.countNonZero(orange_mask)

        # 决策逻辑：只有当一个颜色的像素数明显高于另一个时才做判断，以增加稳定性
        if red_pixel_count > orange_pixel_count * 1.5:
            return "red"
        elif orange_pixel_count > red_pixel_count * 1.5:
            return "orange"
        else:
            return "unknown"  # 如果颜色不明确，则返回未知

    def print_results(self):
        """
        将当前帧的检测、分割和跟踪结果打印到终端，并对红/橙色进行校验。
        """
        if not self.results or not self.results[0] or self.current_frame is None:
            return

        result = self.results[0]
        boxes = result.boxes
        masks = result.masks
        class_names = result.names

        print("--- 实时检测结果 ---")
        if boxes is not None and len(boxes) > 0:
            for i in range(len(boxes)):
                class_id = int(boxes.cls[i].item())
                original_class_name = class_names.get(class_id, "未知类别")
                class_name = original_class_name
                confidence = boxes.conf[i].item()
                
                # 如果是红色或橙色，则启动颜色校验
                refinement_note = ""
                if original_class_name in ["orange", "red"]:
                    bbox = boxes.xyxy[i].cpu().numpy()
                    refined_class = self._refine_red_orange_classification(bbox)
                    if refined_class != "unknown":
                        class_name = refined_class
                        if class_name != original_class_name:
                            refinement_note = f" (修正自: {original_class_name})"

                track_id_info = ""
                if boxes.id is not None:
                    track_id = int(boxes.id[i].item())
                    track_id_info = f", 跟踪ID: {track_id}"

                print(f"  物体 {i+1}: 类别='{class_name}'{refinement_note}, 置信度={confidence:.2f}{track_id_info}")

            if masks is not None and masks.data.numel() > 0:
                print(f"  -> 检测到 {len(masks)} 个物体的实例分割掩码。")
            else:
                print("  -> 未进行实例分割（或模型不支持）。")
        else:
            print("  当前帧未检测到任何物体。")
        
        print("-" * 22 + "\n")

def main():
    """
    主函数，用于启动摄像头并进行实时预测。
    """
    model_path =r"weights/yolo11x-200-epochs/best.pt"
    
    try:
        predictor = YoloPredictor(model_path)
    except Exception as e:
        print(f"错误：无法加载模型，请检查路径 '{model_path}' 是否正确。错误信息: {e}")
        return

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("错误：无法打开摄像头。")
        return

    print("\n摄像头已启动，正在进行实时检测... 按 Ctrl+C 退出。")

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("错误：无法从摄像头捕获帧。")
                break

            predictor.predict_and_track(frame)
            predictor.print_results()

    except KeyboardInterrupt:
        print("\n检测已停止。")
    finally:
        cap.release()
        cv2.destroyAllWindows()
        print("资源已释放。")

if __name__ == "__main__":
    main()
